"""
Caption Generator Module
图像描述生成模块
"""

import logging
from typing import List
from transformers import T5ForConditionalGeneration, T5Tokenizer

logger = logging.getLogger(__name__)


class CaptionGenerator:
    """图像描述生成器"""
    
    def __init__(self, config: dict):
        """
        初始化描述生成器
        
        Args:
            config: 配置字典
        """
        self.config = config
        self.model = None
        self.tokenizer = None
        
        self._load_model()
    
    def _load_model(self):
        """加载FLAN-T5模型和tokenizer"""
        try:
            model_path = self.config['models']['llm_model_path']
            self.model = T5ForConditionalGeneration.from_pretrained(model_path)
            self.tokenizer = T5Tokenizer.from_pretrained(model_path)
            
            # 设置生成参数
            self.max_length = self.config['generation']['max_length']
            self.num_beams = self.config['generation']['num_beams']
            
            logger.info("FLAN-T5模型加载成功")
            
        except Exception as e:
            logger.error(f"FLAN-T5模型加载失败: {e}")
            raise
    
    def build_prompt(self, retrieved_captions: List[List[str]]) -> str:
        """
        构建LLM提示词
        
        Args:
            retrieved_captions: 检索到的图像描述列表
            
        Returns:
            str: 构建的提示词
        """
        try:
            # 提取所有描述
            all_captions = []
            for caption_list in retrieved_captions:
                all_captions.extend(caption_list)
            
            # 去重并限制数量
            unique_captions = list(dict.fromkeys(all_captions))[:10]  # 最多10个描述
            
            # 构建提示词
            descriptions_text = "\n".join([f"- {caption}" for caption in unique_captions])
            
            prompt = f"""基于以下相似图像的描述，请生成一个准确且全面的图像描述：

相似图像描述：
{descriptions_text}

请综合分析以上描述，生成一个新的图像描述："""
            
            return prompt
            
        except Exception as e:
            logger.error(f"构建提示词失败: {e}")
            raise
    
    def generate_caption(self, prompt: str) -> str:
        """
        生成最终描述
        
        Args:
            prompt: 输入提示词
            
        Returns:
            str: 生成的图像描述
        """
        try:
            # 编码输入
            inputs = self.tokenizer(
                prompt,
                return_tensors="pt",
                max_length=512,
                truncation=True,
                padding=True
            )
            
            # 生成描述
            with self.tokenizer.as_target_tokenizer():
                outputs = self.model.generate(
                    inputs.input_ids,
                    max_length=self.max_length,
                    num_beams=self.num_beams,
                    early_stopping=True,
                    do_sample=False,
                    temperature=0.7,
                    no_repeat_ngram_size=2
                )
            
            # 解码输出
            generated_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
            
            # 清理输出文本
            generated_text = generated_text.strip()
            if generated_text.startswith("图像描述："):
                generated_text = generated_text[4:].strip()
            
            logger.info("图像描述生成成功")
            return generated_text
            
        except Exception as e:
            logger.error(f"生成描述失败: {e}")
            raise
